package com.thghh.executor;

import lombok.extern.slf4j.Slf4j;

import java.io.Closeable;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 任务执行的线程池
 *
 * @author pengzhikang
 * @date 2021/9/2 15:37
 */
@Slf4j
    public class TaskPoolExecutor extends ThreadPoolExecutor {

    private static final long TIME_OUT = 60L;
    private static final RejectedExecutionHandler waitPolicy = new WaitPolicy();
    private static final ThreadFactory threadFactory = new DefaultThreadFactory();

    /**
     * 运行中的线程
     */
    private List<Runnable> tasks = new CopyOnWriteArrayList<>();;

    public TaskPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, new CallerRunsPolicy());
    }

    @Override
    public void afterExecute(Runnable r, Throwable t) {
        tasks.remove(r);
    }

    @Override
    public void beforeExecute(Thread t, Runnable r) {
        tasks.add(r);
    }

    public int taskSize() {
        return tasks.size();
    }

    /**
     * 关闭线程池
     *
     * @author zhikang.peng
     * @date 2016年8月4日下午4:28:41
     * @since 1.0
     */
    public void shutdownAndAwaitTermination() {
        shutdown(); // Disable new tasks from being submitted
        try {
            // 通知调度器中所有未完成的任务结束任务，和清空等待队列中的任务
            for (Runnable r : tasks) {
                if (r instanceof Closeable) {
                    ((Closeable) r).close();
                }
            }
            // 清空等待队列
            getQueue().clear();
            // Wait a while for existing tasks to terminate
            if (!awaitTermination(TIME_OUT, TimeUnit.SECONDS)) {
                shutdownNow(); // Cancel currently executing tasks
                // Wait a while for tasks to respond to being cancelled
                if (!awaitTermination(TIME_OUT, TimeUnit.SECONDS)) {
                    log.error("Pool did not terminate");
                }
            }
        } catch (Exception e) {
            // (Re-)Cancel if current thread also interrupted
            shutdownNow();
            // Preserve interrupt status
            Thread.currentThread().interrupt();
        }
    }

    static int size;
    static class DefaultThreadFactory implements ThreadFactory {
        private static final AtomicInteger poolNumber = new AtomicInteger(1);
        private final ThreadGroup group;
        private final AtomicInteger threadNumber = new AtomicInteger(1);
        private final String namePrefix;

        DefaultThreadFactory() {
            SecurityManager s = System.getSecurityManager();
            group = (s != null) ? s.getThreadGroup() :
                    Thread.currentThread().getThreadGroup();
            namePrefix = "task-" +
                    poolNumber.getAndIncrement() +
                    "-thread-";
        }

        @Override
        public Thread newThread(Runnable r) {
            System.out.println("创建了一次线程 " + ++size);
            Thread t = new Thread(group, r,
                    namePrefix + threadNumber.getAndIncrement(),
                    0);
            if (t.isDaemon()) {
                t.setDaemon(false);
            }
            if (t.getPriority() != Thread.NORM_PRIORITY) {
                t.setPriority(Thread.NORM_PRIORITY);
            }
            return t;
        }
    }


    static class WaitPolicy implements RejectedExecutionHandler {

        @Override
        public void rejectedExecution(Runnable task, ThreadPoolExecutor executor) {
            try {
                // 阻塞
                executor.getQueue().put(task);
            } catch (InterruptedException e) {
                log.error("Interrupted ", e);
                // Restore interrupted state...
                Thread.currentThread().interrupt();
            }
        }

    }
}
